import fitz  # PyMuPDF
from typing import Dict, Any, List, Optional
from io import BytesIO
import re

class PDFProcessor:
    """PDF文档处理器"""
    
    def __init__(self):
        self.supported_formats = ['pdf']
    
    def extract_content(self, file_content: bytes) -> Dict[str, Any]:
        """
        从PDF文件中提取内容和结构信息
        
        Args:
            file_content: PDF文件的字节内容
            
        Returns:
            Dict[str, Any]: 包含文本内容和结构信息的字典
        """
        
        try:
            # 从字节流创建PDF文档
            doc = fitz.open(stream=file_content, filetype="pdf")
            
            # 检查文档是否加密
            if doc.needs_pass:
                doc.close()
                raise Exception("PDF文档已加密，无法处理")
            
            # 提取基本信息
            metadata = self._extract_metadata(doc)
            
            # 提取文本内容
            text_content = self._extract_text(doc)
            
            # 提取结构信息
            structure = self._extract_structure(doc)
            
            # 提取表格信息
            tables = self._extract_tables(doc)
            
            # 提取图片信息
            images = self._extract_images(doc)
            
            doc.close()
            
            result = {
                'text': text_content,
                'metadata': metadata,
                'structure': structure,
                'tables': tables,
                'images': images,
                'format': 'pdf'
            }
            
            return result
            
        except Exception as e:
            raise Exception(f"PDF处理失败: {str(e)}")
    
    def _extract_metadata(self, doc: fitz.Document) -> Dict[str, Any]:
        """提取PDF元数据"""
        
        metadata = doc.metadata
        
        return {
            'title': metadata.get('title', ''),
            'author': metadata.get('author', ''),
            'subject': metadata.get('subject', ''),
            'creator': metadata.get('creator', ''),
            'producer': metadata.get('producer', ''),
            'creation_date': metadata.get('creationDate', ''),
            'modification_date': metadata.get('modDate', ''),
            'page_count': doc.page_count,
            'encrypted': doc.needs_pass
        }
    
    def _extract_text(self, doc: fitz.Document) -> str:
        """提取PDF文本内容"""
        
        text_parts = []
        
        for page_num in range(doc.page_count):
            page = doc[page_num]
            
            # 尝试不同的文本提取方法
            text = self._extract_page_text(page, page_num)
            
            if text.strip():
                # 添加页面分隔符
                text_parts.append(f"\n--- 第 {page_num + 1} 页 ---\n")
                text_parts.append(text)
        
        return '\n'.join(text_parts)
    
    def _extract_page_text(self, page: fitz.Page, page_num: int) -> str:
        """提取单页文本"""
        
        try:
            # 方法1：标准文本提取
            text = page.get_text()
            
            # 如果文本为空或过短，尝试其他方法
            if len(text.strip()) < 10:
                # 方法2：按块提取文本
                text = self._extract_text_blocks(page)
            
            # 如果仍然为空，可能是扫描版PDF
            if len(text.strip()) < 10:
                text = f"[第{page_num + 1}页：可能是扫描版，建议使用OCR]"
            
            return text
            
        except Exception as e:
            return f"[第{page_num + 1}页：文本提取失败 - {str(e)}]"
    
    def _extract_text_blocks(self, page: fitz.Page) -> str:
        """按文本块提取内容"""
        
        try:
            blocks = page.get_text("dict")
            text_parts = []
            
            for block in blocks.get("blocks", []):
                if "lines" in block:  # 文本块
                    for line in block["lines"]:
                        line_text = ""
                        for span in line.get("spans", []):
                            line_text += span.get("text", "")
                        
                        if line_text.strip():
                            text_parts.append(line_text)
            
            return '\n'.join(text_parts)
            
        except Exception:
            return ""
    
    def _extract_structure(self, doc: fitz.Document) -> Dict[str, Any]:
        """提取PDF结构信息"""
        
        structure = {
            'headings': [],
            'bookmarks': [],
            'page_info': []
        }
        
        # 提取书签（目录）
        try:
            toc = doc.get_toc()
            for item in toc:
                level, title, page_num = item
                structure['bookmarks'].append({
                    'level': level,
                    'title': title,
                    'page': page_num
                })
        except Exception:
            pass
        
        # 提取页面信息
        for page_num in range(doc.page_count):
            page = doc[page_num]
            page_info = {
                'page_num': page_num + 1,
                'width': page.rect.width,
                'height': page.rect.height,
                'rotation': page.rotation
            }
            structure['page_info'].append(page_info)
        
        # 尝试从文本中识别标题
        structure['headings'] = self._identify_headings_from_text(doc)
        
        return structure
    
    def _identify_headings_from_text(self, doc: fitz.Document) -> List[Dict[str, Any]]:
        """从文本中识别标题"""
        
        headings = []
        
        for page_num in range(doc.page_count):
            page = doc[page_num]
            
            try:
                # 获取文本块信息
                blocks = page.get_text("dict")
                
                for block in blocks.get("blocks", []):
                    if "lines" in block:
                        for line in block["lines"]:
                            line_text = ""
                            font_size = 0
                            font_flags = 0
                            
                            for span in line.get("spans", []):
                                line_text += span.get("text", "")
                                font_size = max(font_size, span.get("size", 0))
                                font_flags |= span.get("flags", 0)
                            
                            line_text = line_text.strip()
                            
                            # 判断是否可能是标题
                            if self._is_likely_heading(line_text, font_size, font_flags):
                                headings.append({
                                    'text': line_text,
                                    'page': page_num + 1,
                                    'font_size': font_size,
                                    'is_bold': bool(font_flags & 2**4),
                                    'level': self._estimate_heading_level(font_size)
                                })
            
            except Exception:
                continue
        
        return headings
    
    def _is_likely_heading(self, text: str, font_size: float, font_flags: int) -> bool:
        """判断是否可能是标题"""
        
        if not text or len(text) > 100:  # 标题通常不会太长
            return False
        
        # 检查字体大小（通常标题字体较大）
        if font_size > 14:
            return True
        
        # 检查是否加粗
        if font_flags & 2**4:  # 加粗标志
            return True
        
        # 检查文本模式
        heading_patterns = [
            r'^第[一二三四五六七八九十\d]+[章节部分]',
            r'^\d+\.\s+\S+',
            r'^[一二三四五六七八九十]、',
            r'^Chapter\s+\d+',
            r'^\d+\s+\S+'
        ]
        
        for pattern in heading_patterns:
            if re.match(pattern, text, re.IGNORECASE):
                return True
        
        return False
    
    def _estimate_heading_level(self, font_size: float) -> int:
        """根据字体大小估计标题级别"""
        
        if font_size >= 20:
            return 1
        elif font_size >= 16:
            return 2
        elif font_size >= 14:
            return 3
        else:
            return 4
    
    def _extract_tables(self, doc: fitz.Document) -> List[Dict[str, Any]]:
        """提取表格信息"""
        
        tables = []
        
        for page_num in range(doc.page_count):
            page = doc[page_num]
            
            try:
                # 尝试查找表格
                page_tables = self._find_tables_in_page(page, page_num)
                tables.extend(page_tables)
            
            except Exception:
                continue
        
        return tables
    
    def _find_tables_in_page(self, page: fitz.Page, page_num: int) -> List[Dict[str, Any]]:
        """在页面中查找表格"""
        
        tables = []
        
        try:
            # 方法1：查找表格边框
            drawings = page.get_drawings()
            
            # 简单的表格检测：查找矩形网格
            rects = []
            for drawing in drawings:
                for item in drawing.get("items", []):
                    if item[0] == "re":  # 矩形
                        rect = item[1]
                        rects.append(rect)
            
            # 如果找到足够的矩形，可能是表格
            if len(rects) > 4:
                # 提取表格区域的文本
                table_text = self._extract_table_text(page, rects)
                if table_text:
                    tables.append({
                        'page': page_num + 1,
                        'content': table_text,
                        'type': 'detected_table',
                        'rect_count': len(rects)
                    })
        
        except Exception:
            pass
        
        return tables
    
    def _extract_table_text(self, page: fitz.Page, rects: List) -> str:
        """从表格区域提取文本"""
        
        try:
            # 找到表格的边界
            if not rects:
                return ""
            
            min_x = min(rect.x0 for rect in rects)
            max_x = max(rect.x1 for rect in rects)
            min_y = min(rect.y0 for rect in rects)
            max_y = max(rect.y1 for rect in rects)
            
            # 创建表格区域
            table_rect = fitz.Rect(min_x, min_y, max_x, max_y)
            
            # 提取该区域的文本
            text = page.get_text(clip=table_rect)
            
            return text
        
        except Exception:
            return ""
    
    def _extract_images(self, doc: fitz.Document) -> List[Dict[str, Any]]:
        """提取图片信息"""
        
        images = []
        
        for page_num in range(doc.page_count):
            page = doc[page_num]
            
            try:
                image_list = page.get_images()
                
                for img_index, img in enumerate(image_list):
                    xref = img[0]
                    
                    # 获取图片信息
                    try:
                        base_image = doc.extract_image(xref)
                        image_info = {
                            'page': page_num + 1,
                            'index': img_index,
                            'xref': xref,
                            'width': base_image.get('width', 0),
                            'height': base_image.get('height', 0),
                            'ext': base_image.get('ext', ''),
                            'size': len(base_image.get('image', b''))
                        }
                        images.append(image_info)
                    
                    except Exception:
                        # 如果无法提取图片，至少记录存在
                        images.append({
                            'page': page_num + 1,
                            'index': img_index,
                            'xref': xref,
                            'error': 'Failed to extract image data'
                        })
            
            except Exception:
                continue
        
        return images
    
    def extract_text_with_coordinates(self, file_content: bytes) -> List[Dict[str, Any]]:
        """提取带坐标信息的文本"""
        
        try:
            doc = fitz.open(stream=file_content, filetype="pdf")
            
            if doc.needs_pass:
                doc.close()
                raise Exception("PDF文档已加密，无法处理")
            
            text_blocks = []
            
            for page_num in range(doc.page_count):
                page = doc[page_num]
                
                blocks = page.get_text("dict")
                
                for block in blocks.get("blocks", []):
                    if "lines" in block:
                        for line in block["lines"]:
                            for span in line.get("spans", []):
                                text = span.get("text", "").strip()
                                if text:
                                    text_blocks.append({
                                        'text': text,
                                        'page': page_num + 1,
                                        'bbox': span.get("bbox", []),
                                        'font': span.get("font", ""),
                                        'size': span.get("size", 0),
                                        'flags': span.get("flags", 0)
                                    })
            
            doc.close()
            return text_blocks
            
        except Exception as e:
            raise Exception(f"坐标文本提取失败: {str(e)}")
    
    def is_text_based_pdf(self, file_content: bytes) -> bool:
        """检查PDF是否基于文本（非扫描版）"""
        
        try:
            doc = fitz.open(stream=file_content, filetype="pdf")
            
            if doc.needs_pass:
                doc.close()
                return False
            
            # 检查前几页的文本内容
            total_text_length = 0
            pages_to_check = min(3, doc.page_count)
            
            for page_num in range(pages_to_check):
                page = doc[page_num]
                text = page.get_text()
                total_text_length += len(text.strip())
            
            doc.close()
            
            # 如果平均每页文本少于50个字符，可能是扫描版
            avg_text_per_page = total_text_length / pages_to_check if pages_to_check > 0 else 0
            
            return avg_text_per_page > 50
            
        except Exception:
            return False